from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import json
import os
import sys
import time

# 添加项目根目录到Python路径，以便能够导入现有的模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from story_generator import StoryGenerator
from image_prompt_generator import ImagePromptGenerator
from utils.file_utils import save_to_file
from storybook_generator import StorybookApp

# 初始化FastAPI应用
app = FastAPI(title="儿童故事绘本创作API", description="基于大模型的儿童故事绘本创作平台")

# 配置CORS，允许前端应用访问API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境中应指定具体的前端域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 加载配置文件
def load_config():
    """加载配置文件"""
    config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'config.json')
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"加载配置文件失败: {e}")
        # 返回默认配置
        return {
            "api_keys": {
                "openai": "your_openai_api_key_here",
                "anthropic": "your_anthropic_api_key_here"
            },
            "story_settings": {
                "max_words": 1000,
                "age_group": "4-8",
                "language": "zh-CN"
            },
            "image_settings": {
                "style": "cartoon",
                "resolution": "1024x1024",
                "number_of_images": 4
            }
        }

# 加载配置
config = load_config()

# 初始化故事生成器和插画提示词生成器
story_generator = StoryGenerator(config)
image_prompt_generator = ImagePromptGenerator(config)
storybook_app = StorybookApp()

# Pydantic模型定义
def create_pydantic_models():
    """动态创建Pydantic模型"""
    class StoryParams(BaseModel):
        theme: str
        main_character: str
        setting: str
        moral: str = ""
        
    class ImagePromptRequest(BaseModel):
        story_content: str
        
    return StoryParams, ImagePromptRequest

StoryParams, ImagePromptRequest = create_pydantic_models()

# API端点定义
@app.get("/")
async def root():
    return {
        "message": "欢迎使用儿童故事绘本创作API",
        "version": "1.0",
        "endpoints": [
            {"path": "/api/story", "method": "POST", "description": "生成儿童故事"},
            {"path": "/api/image-prompts", "method": "POST", "description": "生成插画提示词"},
            {"path": "/api/storybook", "method": "POST", "description": "生成完整故事绘本(故事+插画提示词)"},
            {"path": "/api/storybook-generator", "method": "POST", "description": "使用storybook_generator生成完整故事绘本(故事+插画提示词+图片)"},
            {"path": "/api/output-data", "method": "GET", "description": "获取output目录中的数据"}
        ]
    }

@app.get("/api/output-data")
async def get_output_data():
    """
    获取output目录中的数据
    
    返回:
    - 故事内容
    - 插画提示词列表
    - 是否已生成图片
    """
    try:
        # 构建output目录路径
        output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'output')
        
        # 读取故事内容
        story_file = os.path.join(output_dir, 'story.txt')
        story_content = ''
        if os.path.exists(story_file):
            with open(story_file, 'r', encoding='utf-8') as f:
                story_content = f.read()
        
        # 读取插画提示词
        prompts_file = os.path.join(output_dir, 'image_prompts.txt')
        image_prompts = []
        if os.path.exists(prompts_file):
            with open(prompts_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
                # 解析提示词，格式为 "1. 提示词内容"
                for line in lines:
                    if line.strip() and '.' in line:
                        prompt = line.split('.', 1)[1].strip()
                        image_prompts.append(prompt)
        
        # 检查是否已生成图片
        image_exists = os.path.exists(os.path.join(output_dir, 'story_image.png.png'))
        
        return {
            "success": True,
            "story_content": story_content,
            "image_prompts": image_prompts,
            "image_exists": image_exists,
            "message": "获取output目录数据成功"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取output目录数据失败: {str(e)}")

@app.post("/api/story")
async def generate_story(story_params: StoryParams):
    """
    生成儿童故事
    
    参数:
    - story_params: 故事参数，包含主题、主角、背景等信息
    
    返回:
    - 生成的故事内容
    """
    try:
        # 调用故事生成器生成故事
        story_content = story_generator.generate_story(story_params.dict())
        
        return {
            "success": True,
            "story_content": story_content,
            "message": "故事生成成功"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"生成故事失败: {str(e)}")

@app.post("/api/image-prompts")
async def generate_image_prompts(image_prompt_request: ImagePromptRequest):
    """
    生成插画提示词
    
    参数:
    - image_prompt_request: 包含故事内容的请求
    
    返回:
    - 生成的插画提示词列表
    """
    try:
        # 调用插画提示词生成器生成提示词
        prompts = image_prompt_generator.generate_prompts(image_prompt_request.story_content)
        
        return {
            "success": True,
            "image_prompts": prompts,
            "message": "插画提示词生成成功"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"生成插画提示词失败: {str(e)}")

@app.post("/api/storybook")
async def generate_storybook(story_params: StoryParams):
    """
    生成完整的故事绘本（故事内容+插画提示词）
    
    参数:
    - story_params: 故事参数，包含主题、主角、背景等信息
    
    返回:
    - 生成的故事内容和插画提示词列表
    """
    try:
        # 生成故事
        story_content = story_generator.generate_story(story_params.dict())
        
        # 生成插画提示词
        image_prompts = image_prompt_generator.generate_prompts(story_content)
        
        # 保存结果到文件
        save_to_file(story_content, "story.txt", "output")
        prompts_text = "\n".join([f"{i}. {prompt}" for i, prompt in enumerate(image_prompts, 1)])
        save_to_file(prompts_text, "image_prompts.txt", "output")
        
        return {
            "success": True,
            "story_content": story_content,
            "image_prompts": image_prompts,
            "message": "故事绘本生成成功"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"生成故事绘本失败: {str(e)}")

def cleanup_old_files():
    """
    清理output目录中的旧文件
    删除1小时前生成的图片文件
    """
    try:
        output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'output')
        cutoff_time = time.time() - 3600  # 1小时前的时间戳
        
        # 遍历output目录中的所有文件
        for filename in os.listdir(output_dir):
            file_path = os.path.join(output_dir, filename)
            
            # 只处理文件（不处理目录）
            if os.path.isfile(file_path):
                # 检查文件是否为图片文件
                if filename.endswith('.png') or filename.endswith('.jpg') or filename.endswith('.jpeg'):
                    # 检查文件的修改时间
                    if os.path.getmtime(file_path) < cutoff_time:
                        try:
                            os.remove(file_path)
                            print(f"已删除旧图片文件: {file_path}")
                        except Exception as e:
                            print(f"删除文件时出错: {file_path} - {str(e)}")
    except Exception as e:
        print(f"清理旧文件时出错: {str(e)}")

@app.post("/api/storybook-generator")
async def generate_storybook_with_images(story_params: StoryParams, background_tasks: BackgroundTasks):
    """
    使用storybook_generator生成完整的故事绘本（故事内容+插画提示词+图片）
    
    参数:
    - story_params: 故事参数，包含主题、主角、背景等信息
    
    返回:
    - 生成的故事内容、插画提示词列表和图片信息
    """
    try:
        # 添加后台任务，清理旧文件
        background_tasks.add_task(cleanup_old_files)
        
        # 调用storybook_generator生成完整的故事绘本
        result = storybook_app.generate_storybook(story_params.dict())
        
        return {
            "success": result["success"],
            "story_content": result["story_content"],
            "image_prompts": result["image_prompts"],
            "image_exists": result["image_exists"],
            "image_path": result.get("image_path"),
            "message": "故事绘本（含图片）生成成功"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"生成故事绘本（含图片）失败: {str(e)}")

# 错误处理
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content={
            "success": False,
            "message": f"服务器内部错误: {str(exc)}",
            "error_type": exc.__class__.__name__
        }
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)